# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/29 15:21 
@Author  : ZhangShenao 
@File    : simple_retriever.py 
@Desc    : 自定义检索器

实现自定义的Retriever,需要继承BaseRetriever类,并且实现_get_relevant_documents()方法
异步的方法可以不实现,底层会委托同步方法来执行
"""
from typing import List

from langchain_core.callbacks import CallbackManagerForRetrieverRun
from langchain_core.documents import Document
from langchain_core.retrievers import BaseRetriever


class SimpleRetriever(BaseRetriever):
    """
    自定义检索器
    简单判断检索词query是否存在于Document文档中
    """

    documents: List[Document]
    top_k: int

    def _get_relevant_documents(self, query: str, *, run_manager: CallbackManagerForRetrieverRun) -> List[Document]:
        hit_docs: List[Document] = []

        # 遍历文档列表,如果检索词query存在于文档doc中,则将文档加入hit_docs列表
        for doc in self.documents:
            if len(hit_docs) == self.top_k:
                return hit_docs
            if query.lower() in doc.page_content.lower():
                hit_docs.append(doc)

        # 返回top_k条文档
        return hit_docs
